CycleGAN

CycleGAN

背景介绍

  CycleGAN(Cycle Consistent Generative Adversarial Networks, 循环一致生成式对抗网络):于2017年发表在ICCV上，可以实现图像的风格迁移，风格迁移是GAN网络提出后才出现在人们视野里面的图像处理算法，在生成式对抗网络问世之前，人们很难通过传统的图像处理算法实现风格迁移，今天带小伙伴们看一看瞧一瞧。

CycleGAN理论思想

CycleGAN引入了4个网络结构，分别是生成器GAB，生成器GBA，判别器DA，判别器DB。
GAB的输入是风格A的图像，输出是风格B的图像，目的是将风格A的图像转换为风格B的图像。
GBA的输入是风格B的图像，输出是风格A的图像，目的是将风格B的图像转换为风格A的图像。
DA的输入是风格A的图像，输出是对输入图像的分类，目的是判断输入图像是否为由B转换的风格A的图像
DB的输入是风格B的图像，输出是对输入图像的分类，目的是判断输入图像是否为由A转换的风格B的图像

其中的图像名称有很多，在这里进行简单的介绍。
image_A, image_B指数据集中读取的真实图像，使用DA和DB进行预测时，结果应该是全1。
fake_A指imge_B由GBA生成的风格A类型的图像，使用DA预测时，希望应该是全1，fake_B指imge_A由GAB生成的风格B类型的图像，使用DB预测时，希望应该是全1。
recon_A指fake_B由GBA生成风格A类型的图像，也就是原图image_A经过GAB，再经过GBA生成风格A的图像，希望和image_A越接近越好。
recon_B指fake_A由GAB生成风格B类型的图像，也就是原图image_B经过GBA，再经过GAB生成风格B的图像，希望和image_B越接近越好。
self_A指image_A由GBA生成的风格A类型的图像，因为GBA是将风格B的图像转换为风格A的图像，输入风格A的图像，应该是不会产生变化，希望和image_A越接近越好。
self_B指image_B由GAB生成的风格B类型的图像，因为GAB是将风格A的图像转换为风格B的图像，输入风格B的图像，应该是不会产生变化，希望和image_B越接近越好。

CycelGAN的特点

  使用InstanceNormalization代替BatchNormalization。
  生成器使用下采样+ResNet结构+上采样对图像进行深层特征提取。
  生成器损失函数采用绝对误差，判别器损失函数采用均方误差。
  对生成器损失函数的权重进行调节，使网络更多关注于生成的图像质量。

CycleGAN图像分析

TensorFlow2.0实现

import os
import glob
import numpy as np
import cv2 as cv
from functools import reduce
import tensorflow as tf
import tensorflow.keras as keras


def compose(*funcs):
    if funcs:
        return reduce(lambda f, g: lambda *a, **kw: g(f(*a, **kw)), funcs)
    else:
        raise ValueError('Composition of empty sequence not supported.')


class InstanceNormalization(keras.layers.Layer):
    def __init__(self, beta_initializer='zeros', gamma_initializer='ones',
                 beta_regularizer=None, gamma_regularizer=None,
                 beta_constraint=None, gamma_constraint=None, epsilon=1e-5,
                 **kwargs):
        super(InstanceNormalization, self).__init__(**kwargs)
        self.epsilon = epsilon
        self.beta_initializer = keras.initializers.get(beta_initializer)
        self.gamma_initializer = keras.initializers.get(gamma_initializer)
        self.beta_regularizer = keras.regularizers.get(beta_regularizer)
        self.gamma_regularizer = keras.regularizers.get(gamma_regularizer)
        self.beta_constraint = keras.constraints.get(beta_constraint)
        self.gamma_constraint = keras.constraints.get(gamma_constraint)

    def build(self, input_shape):
        assert len(input_shape) == 4
        self.gamma = self.add_weight(shape=(input_shape[-1],), name='gamma', initializer=self.gamma_initializer,
                                     regularizer=self.gamma_regularizer, constraint=self.gamma_constraint)
        self.beta = self.add_weight(shape=(input_shape[-1],), name='beta', initializer=self.beta_initializer,
                                    regularizer=self.beta_regularizer, constraint=self.beta_constraint)

    def call(self, inputs, **kwargs):
        mean, variance = tf.nn.moments(inputs, axes=[1, 2])
        mean = tf.reshape(mean, shape=[-1, 1, 1, inputs.shape[-1]])
        variance = tf.reshape(variance, shape=[-1, 1, 1, inputs.shape[-1]])
        outputs = (inputs - mean) / tf.sqrt(variance + self.epsilon)
        return outputs * self.gamma + self.beta

    def get_config(self):
        config = {
            'epsilon': self.epsilon,
            'beta_initializer': keras.initializers.serialize(self.beta_initializer),
            'gamma_initializer': keras.initializers.serialize(self.gamma_initializer),
            'beta_regularizer': keras.regularizers.serialize(self.beta_regularizer),
            'gamma_regularizer': keras.regularizers.serialize(self.gamma_regularizer),
            'beta_constraint': keras.constraints.serialize(self.beta_constraint),
            'gamma_constraint': keras.constraints.serialize(self.gamma_constraint)
        }
        base_config = super(InstanceNormalization, self).get_config()

        return dict(list(base_config.items()) + list(config.items()))


class Conv_In_Relu(keras.layers.Layer):
    def __init__(self, filters, kernel_size, strides, padding, name):
        super(Conv_In_Relu, self).__init__()
        self._name = name
        self.block = keras.Sequential([keras.layers.Conv2D(filters, kernel_size, strides, padding)])
        if name.find('in') != -1:
            self.block.add(InstanceNormalization())
        if name.find('leakyrelu') != -1:
            self.block.add(keras.layers.LeakyReLU(0.2))
        elif name.find('relu') != -1:
            self.block.add(keras.layers.ReLU())

    def call(self, inputs, **kwargs):

        return self.block(inputs)


def identity_block(x, filters, kernel_size, strides, padding, name):
    shortcut = x
    x = compose(Conv_In_Relu(filters, kernel_size, strides, padding, name='{}_conv_in_relu1'.format(name)),
                Conv_In_Relu(filters, kernel_size, strides, padding, name='{}_conv_in2'.format(name)))(x)
    x = keras.layers.Add(name='{}_add'.format(name))([x, shortcut])
    x = keras.layers.ReLU(name='{}_relu'.format(name))(x)

    return x


def generator(input_shape, name):
    input_tensor = keras.layers.Input(input_shape, name='input')
    x = input_tensor

    x = compose(Conv_In_Relu(64, (7, 7), (1, 1), 'same', name='conv_in_relu1'),
                Conv_In_Relu(128, (3, 3), (2, 2), 'same', name='conv_in_relu2'),
                Conv_In_Relu(256, (3, 3), (2, 2), 'same', name='conv_in_relu3'))(x)

    for i in range(9):
        x = identity_block(x, 256, (3, 3), (1, 1), 'same', name='identity_block{}'.format(i + 1))

    x = compose(keras.layers.UpSampling2D((2, 2), name='upsampling1'),
                Conv_In_Relu(128, (3, 3), (1, 1), 'same', name='conv_in_relu4'),
                keras.layers.UpSampling2D((2, 2), name='upsampling2'),
                Conv_In_Relu(64, (3, 3), (1, 1), 'same', name='conv_in_relu5'),
                keras.layers.Conv2D(3, (3, 3), (1, 1), 'same', activation='tanh', name='conv_tanh'))(x)

    model = keras.Model(input_tensor, x, name=name)

    return model


def discriminator(input_shape, name):
    input_tensor = keras.layers.Input(input_shape, name='input')
    x = input_tensor

    x = compose(Conv_In_Relu(64, (4, 4), (2, 2), 'same', name='conv_leakyrelu1'),
                Conv_In_Relu(128, (4, 4), (2, 2), 'same', name='conv_in_leakyrelu2'),
                Conv_In_Relu(256, (4, 4), (2, 2), 'same', name='conv_in_leakyrelu3'),
                Conv_In_Relu(512, (4, 4), (2, 2), 'same', name='conv_in_leakyrelu4'),
                keras.layers.Conv2D(1, (3, 3), (1, 1), 'same', name='conv'))(x)

    model = keras.Model(input_tensor, x, name=name)

    return model


def cyclegan(input_shapeA, input_shapeB, model_gAB, model_gBA, model_dA, model_dB):
    input_tensorA = keras.layers.Input(input_shapeA, name='input_A')
    input_tensorB = keras.layers.Input(input_shapeB, name='input_B')

    # 输入风格B由BA生成的风格A类型的图像和输入风格A由AB生成的风格B类型的图像，称为假A和假B
    fake_A = model_gBA(input_tensorB)
    fake_B = model_gAB(input_tensorA)

    # 输入假风格B由BA生成的重建风格A和假风格A由AB生成的重建风格B，称为重建A和重建B
    recon_A = model_gBA(fake_B)
    recon_B = model_gAB(fake_A)

    # 输入风格B由AB生成的风格B类型的图像和输入风格A由BA生成的风格A类型的图像，称为自身B和自身A
    self_A = model_gBA(input_tensorA)
    self_B = model_gAB(input_tensorB)

    model_dA.trainable = False
    model_dB.trainable = False

    conf_A = model_dA(fake_A)
    conf_B = model_dB(fake_B)

    model = keras.Model([input_tensorA, input_tensorB], [conf_A, conf_B, recon_A, recon_B, self_A, self_B], name='CycleGAN')

    return model


def read_data(data_path, img_size, batch_size):
    filename_A = glob.glob(data_path + 'A\\*.jpg')
    filename_B = glob.glob(data_path + 'B\\*.jpg')
    choose_name_A = np.random.choice(filename_A, batch_size)
    choose_name_B = np.random.choice(filename_B, batch_size)

    image_A, image_B = [], []
    for i in range(batch_size):
        A = cv.imread(choose_name_A[i]).astype(np.float32)
        B = cv.imread(choose_name_B[i]).astype(np.float32)
        image_A.append(cv.resize(A, img_size))
        image_B.append(cv.resize(B, img_size))

    image_A = np.array(image_A) / 127.5 - 1
    image_B = np.array(image_B) / 127.5 - 1

    return image_A, image_B


if __name__ == '__main__':
    batch_size = 2
    epochs = 2000
    tf.random.set_seed(22)
    img_size = (128, 128)
    data_path = r'.\monet2photo\train'
    save_path = r'.\cyclegan'
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    optimizer = keras.optimizers.Adam(0.0002, 0.5)
    loss = keras.losses.BinaryCrossentropy()

    real_dAmse = keras.metrics.MeanSquaredError()
    fake_dAmse = keras.metrics.MeanSquaredError()
    real_dBmse = keras.metrics.MeanSquaredError()
    fake_dBmse = keras.metrics.MeanSquaredError()
    gAmse = keras.metrics.MeanSquaredError()
    gBmse = keras.metrics.MeanSquaredError()

    model_dA = discriminator(input_shape=(img_size[0], img_size[1], 3), name='CycleGAN-DiscriminatorA')
    model_dA.compile(optimizer=optimizer, loss='mse')
    model_dB = discriminator(input_shape=(img_size[0], img_size[1], 3), name='CycleGAN-DiscriminatorB')
    model_dB.compile(optimizer=optimizer, loss='mse')

    model_gAB = generator(input_shape=(img_size[0], img_size[1], 3), name='CycleGAN-GeneratorAB')
    model_gBA = generator(input_shape=(img_size[0], img_size[1], 3), name='CycleGAN-GeneratorBA')

    model_gAB.build(input_shape=(img_size[0], img_size[1], 3))
    model_gAB.summary()
    keras.utils.plot_model(model_gAB, 'CycleGAN-generatorAB.png', show_shapes=True, show_layer_names=True)

    model_gBA.build(input_shape=(img_size[0], img_size[1], 3))
    model_gBA.summary()
    keras.utils.plot_model(model_gBA, 'CycleGAN-generatorBA.png', show_shapes=True, show_layer_names=True)

    model_dA.build(input_shape=(img_size[0], img_size[1], 3))
    model_dA.summary()
    keras.utils.plot_model(model_dA, 'CycleGAN-discriminatorA.png', show_shapes=True, show_layer_names=True)

    model_dB.build(input_shape=(img_size[0], img_size[1], 3))
    model_dB.summary()
    keras.utils.plot_model(model_dB, 'CycleGAN-discriminatorB.png', show_shapes=True, show_layer_names=True)

    model = cyclegan(input_shapeA=(img_size[0], img_size[1], 3), input_shapeB=(img_size[0], img_size[1], 3), model_gAB=model_gAB, model_gBA=model_gBA, model_dA=model_dA, model_dB=model_dB)
    model.compile(optimizer=optimizer, loss=['mse', 'mse', 'mae', 'mae', 'mae', 'mae'], loss_weights=[0.5, 0.5, 5, 5, 2.5, 2.5])

    model.build(input_shape=[(img_size[0], img_size[1], 3), (img_size[0], img_size[1], 3)])
    model.summary()
    keras.utils.plot_model(model, 'CycleGAN.png', show_shapes=True, show_layer_names=True)

    for epoch in range(epochs):
        image_A, image_B = read_data(data_path, img_size, batch_size)

        fake_A = model_gBA(image_B)
        fake_B = model_gAB(image_A)

        real_dAmse(np.ones((batch_size, img_size[0] // 16, img_size[1] // 16, 1)), model_dA(image_A))
        fake_dAmse(np.zeros((batch_size, img_size[0] // 16, img_size[1] // 16, 1)), model_dA(fake_A))
        real_dBmse(np.ones((batch_size, img_size[0] // 16, img_size[1] // 16, 1)), model_dB(image_B))
        fake_dBmse(np.zeros((batch_size, img_size[0] // 16, img_size[1] // 16, 1)), model_dB(fake_B))
        gAmse(np.ones((batch_size, img_size[0] // 16, img_size[1] // 16, 1)), model([image_A, image_B])[0])
        gBmse(np.ones((batch_size, img_size[0] // 16, img_size[1] // 16, 1)), model([image_A, image_B])[1])

        real_dAloss = model_dA.train_on_batch(image_A, np.ones((batch_size, img_size[0] // 16, img_size[1] // 16, 1)))
        fake_dAloss = model_dA.train_on_batch(fake_A, np.zeros((batch_size, img_size[0] // 16, img_size[1] // 16, 1)))
        real_dBloss = model_dB.train_on_batch(image_B, np.ones((batch_size, img_size[0] // 16, img_size[1] // 16, 1)))
        fake_dBloss = model_dB.train_on_batch(fake_B, np.zeros((batch_size, img_size[0] // 16, img_size[1] // 16, 1)))

        gloss = model.train_on_batch([image_A, image_B], [np.ones((batch_size, img_size[0] // 16, img_size[1] // 16, 1)), np.ones((batch_size, img_size[0] // 16, img_size[1] // 16, 1)), image_A, image_B, image_A, image_B])

        if epoch % 20 == 0:
            print('epoch = {}, real_dAmse = {}, fake_dAmse = {}, real_dBmse = {}, fake_dBmse = {}, gAmse = {}, gBmse = {}'.format(epoch, real_dAmse.result(), fake_dAmse.result(), real_dBmse.result(), fake_dBmse.result(), gAmse.result(), gBmse.result()))
            real_dAmse.reset_states()
            fake_dAmse.reset_states()
            real_dBmse.reset_states()
            fake_dBmse.reset_states()
            gAmse.reset_states()
            gBmse.reset_states()
            image_A, image_B = read_data(data_path, img_size, batch_size=1)
            fake_A = ((model_gBA(image_B).numpy().squeeze() + 1) * 127.5).astype(np.uint8)
            fake_B = ((model_gAB(image_A).numpy().squeeze() + 1) * 127.5).astype(np.uint8)
            image_A = ((image_A.squeeze() + 1) * 127.5).astype(np.uint8)
            image_B = ((image_B.squeeze() + 1) * 127.5).astype(np.uint8)
            cv.imwrite(save_path + '\\epoch{}.jpg'.format(epoch), np.concatenate([np.concatenate([image_B, fake_A], axis=1), np.concatenate([image_A, fake_B], axis=1)], axis=0))

模型运行结果

小技巧

1. 图像输入可以先将其归一化到0-1之间或者-1-1之间，因为网络的参数一般都比较小，所以归一化后计算方便，收敛较快。
2. 注意其中的一些维度变换和numpy，tensorflow常用操作，否则在阅读代码时可能会产生一些困难。
3. 可以设置一些权重的保存方式，学习率的下降方式和早停方式。
4. CycleGAN对于网络结构，优化器参数，网络层的一些超参数都是非常敏感的，效果不好不容易发现原因，这可能需要较多的工程实践经验。
5. 先创建判别器，然后进行compile，这样判别器就固定了，然后创建生成器时，不要训练判别器，需要将判别器的trainable改成False，此时不会影响之前固定的判别器，这个可以通过模型的_collection_collected_trainable_weights属性查看，如果该属性为空，则模型不训练，否则模型可以训练，compile之后，该属性固定，无论后面如何修改trainable，只要不重新compile，都不影响训练。
6. 在CycleGAN的测试图像中，为了体现模型的效果，第一行的奇数个为拍摄的照片，第一行的偶数个为转换风格后的莫奈风格画作，第二行的奇数个为莫奈风格的画作，第二行的偶数个为转换风格后的照片，这里只是训练了2000代，而且每一代只有2个图像就可以看出CycleGAN的效果。小伙伴们可以选择更大的数据集，更加快速的GPU，训练更长的时间，这样风格迁移的效果就会更加明显。

CycleGAN小结

  CycleGAN是一种有效的风格迁移生成式对抗网络，从上图可以看出CycleGAN模型的参数量只有28M，可以实现任意风格之间的迁移，如果数据集足够，还可以生成人物表情包，是不是非常有趣呢？小伙伴们一定要掌握它。